from flask import Blueprint, request, jsonify
import requests
import json
import os
from backend.app.config import SD_API_URL
from backend.app.utils.http import handle_request_error

bp = Blueprint('sd', __name__)


def load_sd_config(config_name=None):
    config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'config', 'sd_config.json')
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            configs = json.load(f)

            # 如果没有指定配置名称，查找标记为默认的配置
            if config_name is None:
                for name, config in configs.items():
                    if config.get('is_default', False):
                        return config
                # 如果没有找到默认配置，返回名为'default'的配置
                return configs.get('default', {})
            else:
                # 返回指定名称的配置
                return configs.get(config_name, {})
    except Exception as e:
        print(f"Error loading SD config: {e}")
        return {}


@bp.route('/check_sd_status', methods=['GET'])
@handle_request_error
def check_sd_status():
    try:
        response = requests.get(f"{SD_API_URL}/sdapi/v1/sd-models")
        # 包装在 data 中返回
        return jsonify({
            "data": {
                "available": response.status_code == 200
            }
        })
    except:
        return jsonify({
            "data": {
                "available": False
            }
        })


@bp.route('/models', methods=['GET'])
@handle_request_error
def get_models():
    try:
        response = requests.get(f"{SD_API_URL}/sdapi/v1/sd-models")
        if response.status_code == 200:
            models = [model['title'] for model in response.json()]
            return jsonify({
                "data": {
                    "success": True,
                    "models": models
                }
            })
        else:
            return jsonify({
                "data": {
                    "success": False,
                    "error": "获取模型列表失败"
                }
            }), 500
    except Exception as e:
        return jsonify({
            "data": {
                "success": False,
                "error": str(e)
            }
        }), 500


@bp.route('/lora_models', methods=['GET'])
@handle_request_error
def get_lora_models():
    try:
        response = requests.get(f"{SD_API_URL}/sdapi/v1/loras")
        if response.status_code == 200:
            lora_models = [lora['name'] for lora in response.json()]
            return jsonify({
                "data": {
                    "success": True,
                    "lora_models": lora_models
                }
            })
        else:
            return jsonify({
                "data": {
                    "success": False,
                    "error": "获取Lora模型列表失败"
                }
            }), 500
    except Exception as e:
        return jsonify({
            "data": {
                "success": False,
                "error": str(e)
            }
        }), 500


@bp.route('/generate', methods=['POST'])
@handle_request_error
def generate():
    try:
        data = request.json
        config_name = data.get("config_name")
        config = load_sd_config(config_name)

        # 获取批次数
        batch_count = int((config.get("batch_count") if config else None) or 1)

        # 构建请求参数，优先使用用户提供的参数，否则使用加载的配置，并确保数值参数有合适的默认值
        payload = {
            "prompt": data.get("prompt", ""),
            "negative_prompt": data.get("negative_prompt", config.get("negative_prompt", "") if config else ""),
            "width": int(data.get("width") or (config.get("width") if config else None) or 512),
            "height": int(data.get("height") or (config.get("height") if config else None) or 512),
            "steps": int(data.get("steps") or (config.get("steps") if config else None) or 20),
            "cfg_scale": float(data.get("cfg_scale") or (config.get("cfg_scale") if config else None) or 7),
            "sampler_name": data.get("sampler_name", config.get("sampler_name", "Euler a") if config else "Euler a"),
            "seed": int(data.get("seed") or (config.get("seed") if config else None) or -1),
            "enable_hr": bool(data.get("enable_hr") if data.get("enable_hr") is not None else (
                config.get("enable_hr", False) if config else False)),
            "denoising_strength": float(
                data.get("denoising_strength") or (config.get("denoising_strength") if config else None) or 0.7),
            "batch_size": int(data.get("batch_size") or (config.get("batch_size") if config else None) or 1)
        }
        
        # 处理高清修复参数
        hr_resize_mode = data.get("hr_resize_mode", config.get("hr_resize_mode", "scale") if config else "scale")
        
        if hr_resize_mode == "scale":
            # 使用放大倍率模式
            payload["hr_scale"] = float(data.get("upscaler_scale") or (config.get("upscaler_scale") if config else None) or 2)
        else:
            # 使用指定输出尺寸模式
            payload["hr_resize_x"] = int(data.get("hr_target_width") or (config.get("hr_target_width") if config else None) or 1024)
            payload["hr_resize_y"] = int(data.get("hr_target_height") or (config.get("hr_target_height") if config else None) or 1024)

        # 处理Lora模型参数
        lora_model = data.get("lora_model", config.get("lora_model", "") if config else "")
        if lora_model:
            lora_weight = float(data.get("lora_weight") or (config.get("lora_weight") if config else None) or 0.5)
            # 添加Lora模型到提示词中
            payload["prompt"] += f", <lora:{lora_model}:{lora_weight}>"

        # 如果批次数为1，直接生成一次
        if batch_count == 1:
            # 调用SD WebUI API
            response = requests.post(
                f"{SD_API_URL}/sdapi/v1/txt2img",
                json=payload
            )

            if response.status_code == 200:
                result = response.json()
                # 如果batch_size大于1，返回所有生成的图片，否则只返回第一张
                if payload["batch_size"] > 1:
                    return jsonify({
                        "data": {
                            "success": True,
                            "images": result.get("images", [])
                        }
                    })
                else:
                    # 包装在 data 中返回单张图片
                    return jsonify({
                        "data": {
                            "success": True,
                            "image": result.get("images", [])[0] if result.get("images") else None
                        }
                    })
            else:
                return jsonify({
                    "data": {
                        "success": False,
                        "error": "生成失败"
                    }
                }), 500
        else:
            # 如果批次数大于1，生成多批次图片
            images = []
            errors = []

            for i in range(batch_count):
                try:
                    # 每次生成使用不同的种子（如果原始种子为-1）
                    if payload["seed"] == -1:
                        current_payload = {**payload}
                    else:
                        current_payload = {**payload, "seed": payload["seed"] + i if payload["seed"] != -1 else -1}

                    response = requests.post(
                        f"{SD_API_URL}/sdapi/v1/txt2img",
                        json=current_payload
                    )

                    if response.status_code == 200:
                        result = response.json()
                        batch_images = result.get("images", [])
                        images.extend(batch_images)
                    else:
                        errors.append(f"批次 {i + 1} 生成失败")
                except Exception as e:
                    errors.append(f"批次 {i + 1} 生成失败: {str(e)}")

            # 返回所有生成的图片
            return jsonify({
                "data": {
                    "success": len(images) > 0,
                    "images": images,
                    "errors": errors
                }
            })

    except Exception as e:
        return jsonify({
            "data": {
                "success": False,
                "error": str(e)
            }
        }), 500


@bp.route('/batch_generate', methods=['POST'])
@handle_request_error
def batch_generate():
    try:
        data = request.json
        config_name = data.get("config_name")
        config = load_sd_config(config_name)

        # 获取提示词列表
        prompts = data.get("prompts", [])
        if not prompts:
            return jsonify({
                "data": {
                    "success": False,
                    "error": "提示词列表为空"
                }
            }), 400

        # 构建基础请求参数，确保数值参数有合适的默认值并进行类型转换
        base_payload = {
            "negative_prompt": data.get("negative_prompt", config.get("negative_prompt", "") if config else ""),
            "width": int(data.get("width") or (config.get("width") if config else None) or 512),
            "height": int(data.get("height") or (config.get("height") if config else None) or 512),
            "steps": int(data.get("steps") or (config.get("steps") if config else None) or 20),
            "cfg_scale": float(data.get("cfg_scale") or (config.get("cfg_scale") if config else None) or 7),
            "sampler_name": data.get("sampler_name", config.get("sampler_name", "Euler a") if config else "Euler a"),
            "seed": int(data.get("seed") or (config.get("seed") if config else None) or -1),
            "enable_hr": bool(data.get("enable_hr") if data.get("enable_hr") is not None else (
                config.get("enable_hr", False) if config else False)),
            "denoising_strength": float(
                data.get("denoising_strength") or (config.get("denoising_strength") if config else None) or 0.7),
            "hr_scale": float(data.get("upscaler_scale") or (config.get("upscaler_scale") if config else None) or 2)
        }

        # 处理Lora模型参数
        lora_model = data.get("lora_model", config.get("lora_model", "") if config else "")
        lora_weight = float(data.get("lora_weight") or (config.get("lora_weight") if config else None) or 0.5)

        images = []
        errors = []

        # 逐个生成图片
        for prompt in prompts:
            try:
                # 创建带有提示词的payload
                payload = {**base_payload, "prompt": prompt}

                # 如果有Lora模型，添加到提示词中
                if lora_model:
                    payload["prompt"] += f", <lora:{lora_model}:{lora_weight}>"

                response = requests.post(
                    f"{SD_API_URL}/sdapi/v1/txt2img",
                    json=payload
                )

                if response.status_code == 200:
                    result = response.json()
                    image = result.get("images", [])[0] if result.get("images") else None
                    images.append(image)
                else:
                    errors.append(f"生成失败: {prompt}")
                    images.append(None)

            except Exception as e:
                errors.append(f"生成失败: {prompt} - {str(e)}")
                images.append(None)

        return jsonify({
            "data": {
                "success": True,
                "images": images,
                "errors": errors
            }
        })

    except Exception as e:
        return jsonify({
            "data": {
                "success": False,
                "error": str(e)
            }
        }), 500
